#include "kernel_operator.h"

constexpr int32_t TOTAL_LENGTH = 20 * 4;                            // total length of data
constexpr int32_t USE_CORE_NUM = 1;                                   // num of core used
constexpr int32_t BUFFER_NUM = 1;                                     // tensor num for each queue

constexpr int32_t TOTAL_LENGTH_Align32 = (TOTAL_LENGTH + 7) / 8 * 8;
constexpr int32_t SCORES_LENGTH_Align32 = ((TOTAL_LENGTH / 4) + 7)/ 8 * 8;
constexpr int32_t OUT_LENGTH_Align32 = ((TOTAL_LENGTH / 4) + 31)/ 32 * 32;

using namespace AscendC;
class KernelNMSWithMask {
public:
    __aicore__ inline KernelNMSWithMask() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,GM_ADDR z, float iou_threshold, float scores_threshold)  
    {
        this->iou_threshold = iou_threshold;  
        this->scores_threshold = scores_threshold;  

        xGm.SetGlobalBuffer((__gm__ float *)x, TOTAL_LENGTH_Align32);
        yGm.SetGlobalBuffer((__gm__ float *)y, SCORES_LENGTH_Align32);
        zGm.SetGlobalBuffer((__gm__ uint8_t*)z, OUT_LENGTH_Align32);

        pipe.InitBuffer(inQueueX, BUFFER_NUM, TOTAL_LENGTH_Align32 * sizeof(float));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, SCORES_LENGTH_Align32 * sizeof(float));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, OUT_LENGTH_Align32 * sizeof(uint8_t));

        pipe.InitBuffer(tmpBuffer0, OUT_LENGTH_Align32 * sizeof(uint8_t));
    }
    
    __aicore__ inline void Process()
    {
        CopyIn();
        Compute();
        CopyOut();

    }

private:
    __aicore__ inline void CopyIn()
    {
        AscendC::LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        AscendC::LocalTensor<float> yLocal = inQueueY.AllocTensor<float>();
        AscendC::DataCopy(xLocal, xGm[0], TOTAL_LENGTH_Align32);
        AscendC::DataCopy(yLocal, yGm[0], SCORES_LENGTH_Align32);
        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    
    __aicore__ inline void IOU(LocalTensor<float>& xLocal, LocalTensor<float>& yLocal, 
            LocalTensor<uint8_t>& zLocal, LocalTensor<uint8_t>& tmpLocal){
         // -------------- IOU --------------
        // 1. 将所有的边界框按照置信度分数进行降序排序
        // 2. 选择分数最高的边界框，将其加入到保留列表中
        // 3. 计算该边界框与所有剩余边界框的IOU值
        // 4. 删除IOU值高于设定阈值的所有边界框-》label = 0
        // 5. 重复步骤2-4直到所有边界框都被处理
        for(uint32_t i = 0 ; i < (TOTAL_LENGTH / 4) ; i++){
            tmpLocal.SetValue(i, 1); // Initially marked as retain
            zLocal.SetValue(i, 0); // Initialize output mask to 0
        }
         for(uint32_t i = 0 ; i < (TOTAL_LENGTH / 4) ; i++){
             if(tmpLocal.GetValue(i) == 0){
                continue; //该框已被标记为不保留，跳过
            }
            if(yLocal.GetValue(i) < scores_threshold){
                tmpLocal.SetValue(i, 0);
                continue;//低于置信度阈值，跳过
            }
            zLocal.SetValue(i, 1); //标记为保留
            float x1_a = xLocal.GetValue(i * 4 + 0);
            float y1_a = xLocal.GetValue(i * 4 + 1);
            float x2_a = xLocal.GetValue(i * 4 + 2);
            float y2_a = xLocal.GetValue(i * 4 + 3);
            float area1 = (x2_a - x1_a) * (y2_a - y1_a);
            for(uint32_t j = i + 1 ; j < (TOTAL_LENGTH / 4) ; j++){
                if(tmpLocal.GetValue(j) == 0){
                    continue; //该框已被标记为不保留，跳过
                }
                if(yLocal.GetValue(j) < scores_threshold){
                    tmpLocal.SetValue(j, 0);
                    continue;//低于置信度阈值，跳过
                }
                float x1_b = xLocal.GetValue(j * 4 + 0);
                float y1_b = xLocal.GetValue(j * 4 + 1);
                float x2_b = xLocal.GetValue(j * 4 + 2);
                float y2_b = xLocal.GetValue(j * 4 + 3);
                float area2 = (x2_b - x1_b) * (y2_b - y1_b);
                //计算交集坐标
                float x_inter_l = x1_a > x1_b ? x1_a : x1_b;
                float y_inter_l = y1_a > y1_b ? y1_a : y1_b;
                float x_inter_r = x2_a < x2_b ? x2_a : x2_b;
                float y_inter_r = y2_a < y2_b ? y2_a : y2_b;
                
                //计算交集面积
                float w = 0 > (x_inter_r  - x_inter_l) ? 0 :(x_inter_r  - x_inter_l) ;
                float h = 0 > (y_inter_r  - y_inter_l) ? 0 :(y_inter_r  - y_inter_l) ;
                float inter = w * h;

                //计算IOU
                float ovr = inter / (area1 + area2 - inter);
                if(ovr > this->iou_threshold){
                    tmpLocal.SetValue(j, 0); // Mark as discarded
                }
            }
        }
    }

    __aicore__ inline void Compute()
    {

        //========================Preprocess========================
        AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        AscendC::LocalTensor<float> yLocal = inQueueY.DeQue<float>();
        AscendC::LocalTensor<uint8_t> zLocal = outQueueZ.AllocTensor<uint8_t>();
        AscendC::LocalTensor<uint8_t> tmpLocal = tmpBuffer0.Get<uint8_t>();

        //========================Process========================
        IOU(xLocal, yLocal, zLocal, tmpLocal);
        
        //========================Postprocess========================
        outQueueZ.EnQue<uint8_t>(zLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
    }
    
    __aicore__ inline void CopyOut()
    {
        AscendC::LocalTensor<uint8_t> zLocal = outQueueZ.DeQue<uint8_t>();
        AscendC::DataCopy(zGm[0], zLocal, OUT_LENGTH_Align32);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX,inQueueY;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> tmpBuffer0;
    AscendC::GlobalTensor<float> xGm;
    AscendC::GlobalTensor<float> yGm;
    AscendC::GlobalTensor<uint8_t> zGm;
    float iou_threshold;  
    float scores_threshold;
};

extern "C" __global__ __aicore__ void nms_with_mask_custom(GM_ADDR x,GM_ADDR y, GM_ADDR z, float iou_threshold, float scores_threshold)
{
    KernelNMSWithMask op;
    op.Init(x,y, z, iou_threshold, scores_threshold);  
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG
void nms_with_mask_custom_do(uint32_t blockDim, void *stream, uint8_t *x, uint8_t *y,uint8_t *z, float iou_threshold, float scores_threshold)
{
    nms_with_mask_custom<<<blockDim, nullptr, stream>>>(x, y,z, iou_threshold, scores_threshold);
}
#endif